package org.example;

/**
 * 二叉搜索树实现
 * @param <K>
 * @param <V>
 */
public class BSTree<K extends Comparable<K>, V> implements ITree<K, V>{
    static class BSNode<K, V> implements ITreeNode<K, V>{
        private ITreeNode<K, V> parent;
        private ITreeNode<K, V> left;
        private ITreeNode<K, V> right;

        private K key;
        private V value;

        public BSNode(BSNode<K, V> parent, K key, V value) {
            this.parent = parent;
            this.key = key;
            this.value = value;
        }

        public ITreeNode<K, V> getParent() {
            return parent;
        }

        @Override
        public void setParent(ITreeNode<K, V> parent) {
            this.parent = parent;
        }

        @Override
        public ITreeNode<K, V> getLeft() {
            return left;
        }

        @Override
        public void setLeft(ITreeNode<K, V> left) {
            this.left = left;
        }

        @Override
        public ITreeNode<K, V> getRight() {
            return right;
        }

        @Override
        public void setRight(ITreeNode<K, V> right) {
            this.right = right;
        }

        @Override
        public K getKey() {
            return key;
        }

        @Override
        public void setKey(K key) {
            this.key = key;
        }

        @Override
        public V getValue() {
            return value;
        }

        @Override
        public void setValue(V value) {
            this.value = value;
        }
    }

    private ITreeNode<K, V> root;

    public ITreeNode<K, V> getRoot() {
        return root;
    }

    public void setRoot(ITreeNode<K, V> root) {
        this.root = root;
    }

    @Override
    public void put(K key, V value) {
        put(new BSNode<>(null, key, value));
    }

    protected void put(ITreeNode<K, V> BSNode) {
        K key = BSNode.getKey();
        if (key == null) { throw new NullPointerException("Key could not be null"); }
        if (root == null) {
            root = BSNode;
            return;
        }
        ITreeNode<K, V> temp = root;
        ITreeNode<K, V> parent;
        int cmp;
        do{
            parent = temp;
            cmp = key.compareTo(temp.getKey());
            if (cmp > 0) {
                temp = temp.getRight();
            } else if (cmp < 0) {
                temp = temp.getLeft();
            } else {
                temp.setValue(BSNode.getValue());
                return;
            }
        } while (temp != null);

        if (cmp > 0) {
            parent.setRight(BSNode);
        } else {
            parent.setLeft(BSNode);
        }
        BSNode.setParent(parent);
    }

    protected ITreeNode<K, V> getNode(K key) {
        if (key == null) { throw new NullPointerException("Key could not be null"); }
        if (root != null) {
            ITreeNode<K, V> temp = root;
            do {
                int cmp = key.compareTo(temp.getKey());
                if (cmp > 0) {
                    temp = temp.getRight();
                } else if (cmp < 0) {
                    temp = temp.getLeft();
                } else {
                    return temp;
                }
            } while(temp != null);
        }
        return null;
    }

    @Override
    public V get(K key) {
        ITreeNode<K, V> ret = getNode(key);
        return ret == null ? null : ret.getValue();
    }

    protected ITreeNode<K, V> getPredecessor(ITreeNode<K, V> node) {
        if (node == null) { return null; }
        node = node.getLeft();
        ITreeNode<K, V> ret = null;
        while (node != null) {
            ret = node;
            node = node.getRight();
        }
        return ret;
    }

    private void switchNode(ITreeNode<K, V> node1, ITreeNode<K, V> node2) {
        K tempK = node1.getKey();
        V tempV = node1.getValue();
        node1.setKey(node2.getKey());
        node1.setValue(node2.getValue());

        node2.setKey(tempK);
        node2.setValue(tempV);
    }

    @Override
    public V remove(K key) {
        ITreeNode<K, V> temp = getNode(key);
        if (temp == null) {
            return null;
        }
        ITreeNode<K, V> predecessor = getPredecessor(temp);
        if (predecessor != null) {
            switchNode(temp, predecessor);
            if (predecessor.getParent().getRight() == predecessor) {
                predecessor.getParent().setRight(predecessor.getLeft());
            } else {
                predecessor.getParent().setLeft(predecessor.getLeft());
            }
            predecessor.setParent(null);
            return predecessor.getValue();
        }
        if (temp.getParent().getRight() == temp) {
            temp.getParent().setRight(temp.getRight());
        } else {
            temp.getParent().setLeft(temp.getRight());
        }
        temp.setParent(null);
        return temp.getValue();
    }
}
